import cv2
import os
import matplotlib.pyplot as plt
import imageio


def nii_mask_dcm():
    nii_folder = 'nii_folder/'
    dcm_folder = 'dcm_folder/'
    both_next_path = ['9614225_1/', '9614225_2/', '9625145_1/', '9625145_2/']
    for i in range(4):
        dcm_folder_ = os.path.join(dcm_folder, both_next_path[i])
        nii_folder_ = os.path.join(nii_folder, both_next_path[i])
        for root, ds, fs in os.walk(dcm_folder_):
            idx = 0
            for f in fs:
                dcm_path = os.path.join(dcm_folder_, f)
                nii_path = os.path.join(nii_folder_, f)
                img = cv2.imread(dcm_path, 0)
                mask = cv2.imread(nii_path, 0)
                img_masked = cv2.bitwise_and(img, img, mask=mask)
                save_masked_path = os.path.join('./nii_mask_dcm/', both_next_path[i])
                if not os.path.exists(save_masked_path):
                    os.mkdir(save_masked_path)
                plt.figure(figsize=(8, 8))
                plt.imshow(img_masked, 'gray')
                plt.xticks([])
                plt.yticks([])
                plt.axis('off')
                plt.savefig(os.path.join(save_masked_path, '{}.png'.format(idx)), dpi=64)
                idx += 1
                plt.close()


if __name__ == '__main__':
    nii_mask_dcm()